import os
from sklearn.metrics import  recall_score, roc_auc_score, accuracy_score, confusion_matrix
from util import *
from sklearn.metrics import roc_curve, auc
from keras.callbacks import  ModelCheckpoint
import scipy.misc as mc

data_location = ''
#定义测试数据的路径
testing_images_loc = data_location + 'Crack/test/image/'
testing_label_loc = data_location + 'Crack/test/label/'
#获取测试数据文件列表
test_files = os.listdir(testing_images_loc)
test_data = []  #用于存储测试图像数据
test_label = []  #用于存储测试标签数据
desired_size=480
#处理测试数据
for i in test_files:
    im = mc.imread(testing_images_loc + i)
    label = mc.imread(testing_label_loc + "Image_" + i.split('_')[1].split(".")[0] + "_1stHO.png")
    old_size = im.shape[:2]
    delta_w = desired_size - old_size[1]
    delta_h = desired_size - old_size[0]

    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
    left, right = delta_w // 2, delta_w - (delta_w // 2)

    color = [0, 0, 0]
    color2 = [0]

    new_im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                value=color)

    new_label = cv2.copyMakeBorder(label, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                   value=color2)

    test_data.append(cv2.resize(new_im, (desired_size, desired_size)))
    temp = cv2.resize(new_label,
                      (desired_size, desired_size))
    _, temp = cv2.threshold(temp, 127, 255, cv2.THRESH_BINARY)
    test_label.append(temp)
#将测试数据转换为numpy数组
test_data = np.array(test_data)
test_label = np.array(test_label)
#归一化数据
x_test = test_data.astype('float32') / 255.

y_test = test_label.astype('float32') / 255.
#调整数据形状适应模型输入
x_test = np.reshape(x_test, (len(x_test), desired_size, desired_size, 3))
y_test = np.reshape(y_test, (len(y_test), desired_size, desired_size, 1))
#裁剪以适应模型的输入形状
y_test=crop_to_shape(y_test,(len(y_test), 320, 480, 1))

from  RSAN import *
model=RSANet(input_size=(desired_size,desired_size,3),start_neurons=16,keep_prob=0.78,lr=1e-3)
weight="Crack/Model/RSAN.h5"

if os.path.isfile(weight): model.load_weights(weight)

model_checkpoint = ModelCheckpoint(weight, monitor='val_acc', verbose=1, save_best_only=True)
#进行预测
y_pred = model.predict(x_test)
#裁剪测试结果以适应目标形状
y_pred= crop_to_shape(y_pred,(6,320,480,1))
#二值化处理预测结果保存结果图像
y_pred_threshold = []
i=0
for y in y_pred:

    _, temp = cv2.threshold(y, 0.5, 1, cv2.THRESH_BINARY)
    y_pred_threshold.append(temp)
    y = y * 255
    cv2.imwrite('./Crack/test/result/%d.png' % i, y)  #保存结果图像
    i+=1
#将预测结果和标签展平以计算评估指标
y_test = list(np.ravel(y_test))
y_pred_threshold = list(np.ravel(y_pred_threshold))
#计算混淆矩阵
tn, fp, fn, tp = confusion_matrix(y_test, y_pred_threshold).ravel()
#输出评价指标
print('Sensitivity:', recall_score(y_test, y_pred_threshold))
print('Specificity:', tn / (tn + fp))
print("F1:",2*tp/(2*tp+fn+fp))
print('Accuracy:', accuracy_score(y_test, y_pred_threshold))
print('AUC:', roc_auc_score(y_test, list(np.ravel(y_pred))))

